from config import *
import numpy as np
# import numba
import pickle

def create_seed():
    l=num_inputs+num_outputs+num_hidden
    nodes=np.random.randn(l)
    links=np.random.randn(l,l)
    # link不作用于自己
    for i in range(l):
        links[i][i]=0
    # inputs不受作用
    for link in links:
        for i in range(num_inputs):
            link[i]=0
    
    seed=[nodes,links]
    return(seed)

class Brain:
    def __init__(self,seed:list) -> None:
        self.seed=seed
        self.nodes=seed[0]
        self.links=seed[1]

    # @numba.jit(nopython=True)
    def activate(self):
        # 传输
        for i,node in enumerate(self.nodes):
            if node>0:
                self.nodes=self.nodes+node*self.links[i]
                self.links[i]=self.links[i]*link_happy_rate
            else:
                self.links[i]=self.links[i]*link_nohappy_rate

        # 衰弱
        self.nodes=self.nodes*node_weak_rate


    def think(self,inputs):
        for i,input in enumerate(inputs):
            self.nodes[i]=input
        self.activate()

        output=[]
        s=len(inputs)
        for i in range(num_outputs):
            output.append(self.nodes[s+i])
        
        return output

def train(eval,seeds_data):
    # 测试
    for seed_data in seeds_data:
        if seed_data[1]:
            continue
        else:
            br=Brain(seed_data[0])
            seed_data[1]=eval(br.think)

    # 按fitness排序，降序
    seeds_data.sort(reverse=True,key=lambda x:x[1])

    # 变异
    winner_num=int(train_size/3)
    new_seeds_data=[]
    for i in range(winner_num):
        seed_data=seeds_data[i]
        new_seeds_data.append(seed_data)

        for _ in range(2):
            new_seed=seed_data[0].copy()
            shape=new_seed[0].shape
            # new_seed[0]=new_seed[0]*(1+np.random.randn(shape[0],)/10)
            new_seed[0]=new_seed[0]+np.random.randn(shape[0],)/10
            shape=new_seed[1].shape
            new_seed[1]=new_seed[1]+np.random.randn(shape[0],shape[1])/10

            new_seeds_data.append([new_seed,0])

    return new_seeds_data

def retrain(eval,train_times,seeds_data):
    for i in range(train_times):
        seeds_data=train(eval,seeds_data)
        print('第%d次训练，最好fitness：%f'%(i,seeds_data[0][1]))

    return seeds_data

def ran_seeds_data():
    # [seed,fitness]
    seeds_data=[]
    for i in range(train_size):
        seeds_data.append([create_seed(),0])
    return seeds_data

def save_seeds_data(seeds_data):
    with open('seeds_data.pickle','wb') as f:
        pickle.dump(seeds_data,f)

def read_seeds_data():
    with open('seeds_data.pickle','rb') as f:
        seeds_data=pickle.load(f)

    return seeds_data


# #%%
# a=[1,2]
# print(a)
# def f(a):
#     a[0]=11
# f(a)
# print(a)


if __name__=='__main__':
    seed=create_seed()
    print(seed[0])
    print(seed[1])
# %%
